from typing import Dict, Any
import cv2
import numpy as np
from paddleocr import PPStructureV3
from app.utils.exceptions import ModelInitializationError, ImageProcessingError
from app.utils.logger import get_logger
from app.utils.timing import timing_decorator

logger = get_logger(__name__)


# 化版面分析模块

class LayoutAnalysisModule:
    _instance = None
    _initialized = False

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(LayoutAnalysisModule, cls).__new__(cls)
        return cls._instance

    def __init__(self):
        if not self._initialized:
            logger.info("Initializing LayoutAnalysisModule")
            try:
                # 使用PP-Structure进行版面分析
                self.layout_analyzer = PPStructureV3()
                self._initialized = True
            except Exception as e:
                raise ModelInitializationError(f"Failed to initialize layout analysis model: {str(e)}")

    @timing_decorator()
    def analyze(
            self,
            image_data: bytes,
            detect_tables: bool = True,
            detect_formulas: bool = True
    ) -> Dict[str, Any]:
        """执行版面分析"""
        try:
            # 将字节数据转换为OpenCV图像
            nparr = np.frombuffer(image_data, np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

            if img is None:
                raise ImageProcessingError("Failed to decode image data")

            height, width = img.shape[:2]

            # 使用最新版本的 PP-Structure API
            result = self.layout_analyzer(img)  # 修改这一行

            # 处理版面元素
            layout_elements = []
            for element in result:
                layout_element = {
                    "type": element['type'],
                    "bbox": element['box'],
                    "confidence": element['score']
                }

                # 根据类型添加额外信息
                if element['type'] == 'text':
                    layout_element['text'] = element.get('text', '')
                elif element['type'] == 'table':
                    layout_element['html'] = element.get('html', '')

                layout_elements.append(layout_element)

            return {
                "layout_elements": layout_elements,
                "image_size": {"width": width, "height": height}
            }

        except Exception as e:
            logger.exception(f"Layout analysis failed: {str(e)}")
            raise ImageProcessingError(f"Layout analysis failed: {str(e)}")
